import taichi as ti

ti.init(arch=ti.cuda)

n = 180
pixels = ti.Vector.field(3, ti.f32, shape=(n * 2, n, 5 * n))


@ti.func
def complex_sqr(z):
    return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])


@ti.kernel
def paint(t: float):
    for i, j, k in pixels:
        # c = ti.Vector([-0.8,  ti.cos(t) * 0.2])
        # z = ti.Vector([i / n - 1, j / n - 0.5]) * 3 
        c = ti.Vector([
            i / n - 0.5*2, 
            j / n - 0.5*1,
            k / n - 0.5*5]) * 3
        z = ti.Vector([0.0, (1 - ti.cos(t)) * 0.5,0.0])
        iterations = 0
        while z.norm() < 5 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j, k] = [20, 1 - iterations * 0.02, 90]


gui = ti.GUI("mandelbrot set ", pixels.shape)

for i in range(1000000):
    paint(i * 0.03)
    gui.set_image(pixels)
    gui.show()
 
window = ti.ui.Window("mandelbrot set 3D", pixels, vsync=True)
canvas = window.get_canvas()

frame_id = 0

while window.running:
    frame_id += 1 

    for i in range(1000000):
        paint(i * 0.03)
        canvas.set_image(pixels)
        # gui.show()
        # canvas.set_image(julia.get_image(frame_id / 60)) 
    window.show()
